%% Use OptDMD method to extract coherent strcutre from the secondary flow
% Sijie Sun et. cl. PNAS, 2023
% Based on Travis Askham 2017, OptDMD method
%% Choose Dataset
close all;clc;clear
load('Velocity.mat')
addpath('./src');
%% Setup
nt=size(U_tensor,1);%number of frames
nx=size(U_tensor,2);%x dimension
ny=size(U_tensor,3);%y dimension
dt=t(2)-t(1);
mm=1e-3;%1 mm = 1e-3 m
%% Preview Velocity Field
figure()    
quiver(X/mm,Y/mm,squeeze(mean(U_tensor(:,:,:),1)),squeeze(mean(V_tensor(:,:,:),1)),2,'LineWidth',1.5)
axis equal
xlabel('x (mm)')
ylabel('y (mm)')

%%
% v = VideoWriter('Velocity Field.mp4','MPEG-4');
% v.FrameRate=20;v.Quality=100;
% open(v);
% fig=figure('units','pixels','position',[0 0 900 900]);
% for i=1:nt/2
%     U_plot=squeeze(U_tensor(i,:,:))/mm;
%     V_plot=squeeze(V_tensor(i,:,:))/mm;
%     U_plot(1,35)=30;
%     a=quiver(X/mm,Y/mm,U_plot,V_plot,2,'LineWidth',1.5);
%     a.Color=[0.2118    0.2706    0.3098];
%     set(fig,'color','w');
%     title(['t = ' num2str(t(i),'%5.2f') ' s'],'FontSize',20)
%     text(44.2,2.7, '30 mm/s', 'HorizontalAlignment', 'center','FontSize',12);
%     axis equal
%     axis off
%     drawnow
%     writeVideo(v,frame2im(getframe(fig)))
% end
% close(v);
% disp('COMPLETED');
% %% Merge 2D Velocity field into one vector
u=zeros(nx*ny*2,nt);
for i=1:nt
    u(:,i)=UV2X(reshape(U_tensor(i,:,:),[],1),reshape(V_tensor(i,:,:),[],1));
end
%% SVD decomposition with truncation for reference
[U, S, V]=svd(u,'econ');
figure()
stem(diag(S).^2);% Importances of the first ten SVD modes
set(gca,'YScale','log')
xlim([0,10])
% truncation at r
r =21;
U = U(:,1:r);
S = S(1:r,1:r);
V = V(:,1:r);
uhat=U*S*V';
figure()
hold on
plot(t,mean(u,1),'r-','DisplayName','Ground Truth')
plot(t,mean(uhat,1),'b--','linewidth',2,'DisplayName','SVD with truncation')
xlim([0,30])
xlabel('Time-s')
ylabel('Speed-m/s')

legend('Location','best')
set(gca,'fontsize',16)

%% Optimized DMD, On truncated modes

r=21;
OPT.maxiter=2000;
[Phi_optDMD,eigs_optDMD,b_optDMD]=optdmd(u, t,r,2,OPT);
uxhat2=zeros(1,nt);
ux_true=zeros(1,nt);
for j=1:nt
    reconstructed=zeros(size(u(:,1)));reconstructed2=reconstructed;
    for i=1:r
        reconstructed=reconstructed+b_optDMD(i)*exp(eigs_optDMD(i)*t(j))*Phi_optDMD(:,i);
    end     
    uxhat2(j)=mean(real(reconstructed(1:2:end)));
    ux_true(j)=mean(u(1:2:end,j));
end

%% Compare decomposited flow with the measured flow
f=figure('Units',"inches",'Position',[1,1,1.8,1.4317]);

hold on
a=plot(t,ux_true/mm,'DisplayName','Measured Value','linewidth',2,'Color','#377eb8');
c=plot(t,uxhat2/mm,'DisplayName','Reconstruction Value','linewidth',2,'Color','#4daf4a');
a.Color=[a.Color,0.5];
c.Color=[c.Color,0.5];

%%
hold off
legend('Location','best')
legend 'boxoff'
% set(gca,'fontsize',8)
% xlabel('Time (s)')
% ylabel('Ux (mm/s)')
xlim([0,30])
ylim([-4,4])
box on
drawnow() 
%%
set(gca,'linewidth',1)
set(gca,'fontname','Arial') 
% ax=gca;
% ax.Position = [0 0 1 1];
%%
saveas(gcf,['DMD reconstruction'],'epsc')
print(gcf,['DMD reconstruction Comparison'],'-dpng','-r500')
%% Check L2 Difference
L2difference=0;
L2true=0;
for j=1:nt
    reconstructed=zeros(size(u(:,1)));
    for i=1:r
        reconstructed=reconstructed+b_optDMD(i)*exp(eigs_optDMD(i)*dt*j)*Phi_optDMD(:,i);
    end    
    L2difference=L2difference+norm(u(:,j)-reconstructed);
    L2true=L2true+norm(u(:,j));
end
disp(['relative L2 difference is ',num2str(L2difference/L2true),'L2 original is ',num2str(L2true)])

%% Plot  DMD spectrum
f2=figure();
hold on

theta = (0:1:100)*2*pi/100;
plot(cos(theta),sin(theta),'k--') % plot unit circle
scatter(real(exp(eigs_optDMD*dt)),imag(exp(eigs_optDMD*dt)),'ok')
axis([-1.1 1.1 -1.1 1.1]);
xlabel('Re of Eigen Value')
ylabel('Im of Eigen Value')
title('DMD Eigen Value Distribution')
axis equal
print(f2,['Spectrum of Modes'],'-dpng','-r500')

%% Plot DMD Mode Importance Distribution
f3=figure();
[Contribution,Order]=sort(abs(b_optDMD),'descend');
bar(Contribution(1:11)/sum(Contribution),'linewidth',0.0001)
ylim([5e-3,1])
set(gca,'yscale','log')
box on; grid on
xlabel('Mode Number')
ylabel('Contribution')
title('Mode Contribution Distribution')
%% Plot Flow structure Importance Distribution
Y1=[0;Contribution];Y2=zeros(11,0);
for i=1:11
    Y2(i)=sum(Y1(i*2-1:i*2));
end
f4=figure('unit','inches',"Position",[1,1,1.5,1.2]);
bars=bar(Y2(1:6)/sum(Y2),'linewidth',0.0001);
bars.FaceColor=[0, 51, 102]/255;
ylim([5e-3,1])
set(gca,'yscale','log')
box on; grid off
xlabel('Structure Number')
ylabel('Contribution')
set(gca,'Position',[0,0,1,1])

set(gca,'fontsize',10)
set(gca,'linewidth',2)
print(f4,['Structure Contribution Distribution'],'-dpng','-r500')
saveas(f4,['Structure Contribution Distribution'],'epsc')

%% Plot the eigen modes
Scale=2;
for j=1:2:9
    i=Order(j);
    [f,ax1,ax2]=plotVectorfield(Phi_optDMD(:,i),nx,ny,exp(eigs_optDMD(i)*dt),dt,Scale,X,Y,mm);
    print(f,['Velocity field of Sorted Mode ' num2str(j)],'-dpng','-r300')
    saveas(f,['Velocity field of Sorted Mode ' num2str(j)],'epsc')
%     Generate_GIF(Phi_optDMD(:,i),nx,ny,exp(eigs_optDMD(i)*dt),dt,['Mode',num2str(j)]);
end

%% Auxiliary Functions

function X=UV2X(U,V)
    X=zeros(length(U)*2,1);
    X(1:2:end-1)=U;
    X(2:2:end)=V;
end

function [f,ax1,ax2]=plotVectorfield(u,nx,ny,eig_val,dt,Scale,X,Y,mm)
    Uxr=reshape(real(u(1:2:end-1)),nx,ny);
    Uyr=reshape(real(u(2:2:end)),nx,ny);
    Uxi=reshape(imag(u(1:2:end-1)),nx,ny);
    Uyi=reshape(imag(u(2:2:end)),nx,ny);    
    Lr=norm(real(u(:)));
    Li=norm(imag(u(:)));
    f=figure('unit','inches','Position',[2,2,8,3.5]);
    ax1=subplot(1,2,1);
        a=quiver(X/mm,Y/mm,Uxr/Lr,Uyr/Lr,Scale,'LineWidth',0.75)
        a.Color=[99, 109, 111]/255;
        axis equal
        title('Real Part')
        xlabel('x (mm)')
        ylabel('y (mm)')
        set(gca,'linewidth',2)
        set(gca,'fontsize',12)    
        axis off

    ax2=subplot(1,2,2);
        a=quiver(X/mm,Y/mm,Uxi/Li,Uyi/Li,Scale,'LineWidth',0.75)
        a.Color=[99, 109, 111]/255;
        axis equal
        xlabel('x (mm)')
        ylabel('y (mm)')
        set(gca,'linewidth',2)
        set(gca,'fontsize',12)
        title('Imaginary Part')
    dtheta=angle(eig_val);
    omega=dtheta/dt;
    tau=2*pi/omega;
    sgtitle(['% \lambda = ',num2str(eig_val,4) ,' \tau =' ,num2str(abs(tau),4),' s'],'fontsize',16) 
    axis off
end


